import torch
import numpy as np


def sigmoid_np(x):
    return 1.0 / (1.0 + np.exp(-x))


def decode_location2(ploc, row, col, cntK, cntB):
    # ploc: [num_obj, 4]
    # pbbox: [num_obj, 4]
    # cntK, cntB: y = k*x + b
    if isinstance(ploc, torch.Tensor):
        ploc = ploc.t()
        sigmoid = torch.sigmoid
        exp = torch.exp
        stack = torch.stack
    elif isinstance(ploc, np.ndarray):
        ploc = ploc.transpose()
        sigmoid = sigmoid_np
        exp = np.exp
        stack = np.stack
    else:
        raise TypeError("Error data type:", type(ploc), "unsupported for method xywh2xyxy.")

    x_offset, y_offset, w_scale, h_scale = ploc
    cx, cy, w, h = (
        sigmoid(x_offset) + col,
        sigmoid(y_offset) + row,
        exp(sigmoid(w_scale) * cntK + cntB),
        exp(sigmoid(h_scale) * cntK + cntB),
    )
    pbbox = stack([cx, cy, w, h], 1)
    return pbbox


def decode_bbox(bbox, row, col, stride, k, f_stack):
    cx = ((bbox[:, 0] + 1) / 2 + col) * stride
    cy = ((bbox[:, 1] + 1) / 2 + row) * stride
    w = (bbox[:, 2] + 1) * k
    h = (bbox[:, 3] + 1) * k
    bbox = f_stack([cx, cy, w, h], 1)
    return bbox


def ious(box0, box1):
    if len(box0.shape) == 1:
        box0 = box0.unsqueeze(0)
    if len(box1.shape) == 1:
        box1 = box1.unsqueeze(0)
    size_0 = box0.size(0)
    size_1 = box1.size(0)

    # overlaps
    max_xy = torch.min(
        box0[:, 2:].unsqueeze(1).expand(size_0, size_1, 2),
        box1[:, 2:].unsqueeze(0).expand(size_0, size_1, 2),
    )
    min_xy = torch.max(
        box0[:, :2].unsqueeze(1).expand(size_0, size_1, 2),
        box1[:, :2].unsqueeze(0).expand(size_0, size_1, 2),
    )
    inter = torch.clamp((max_xy - min_xy), min=0)
    inters = inter[:, :, 0] * inter[:, :, 1]

    # union
    area_0 = (box0[:, 2] - box0[:, 0]) * (box0[:, 3] - box0[:, 1])
    area_0 = area_0.unsqueeze(1).expand(size_0, size_1)
    area_1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
    area_1 = area_1.unsqueeze(0).expand(size_0, size_1)
    unions = area_0 + area_1 - inters

    # iou
    eps = 1e-7
    ious = inters / (unions + eps)
    return ious
